import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import cv2
import os
import denoise
import numpy as np


class DenoiseApp:
    """
    图像去噪图形界面应用，基于近似最近邻搜索与低秩逼近算法。
    """

    def __init__(self, root):
        """
        初始化界面窗口、变量与按钮控件。

        :param root: 主窗口对象（Tk实例）
        """
        self.root = root
        self.root.title("基于近似最近邻搜索与低秩逼近的非局部图像去噪")
        self.root.geometry("1200x600")

        self.image_path = None
        self.original_img = None
        self.result_img = None

        self.btn_select = tk.Button(root, text="选择图像", command=self.select_image, font=("Arial", 14))
        self.btn_select.pack(pady=10)

        self.btn_process = tk.Button(root, text="开始处理", command=self.process_image, font=("Arial", 14))
        self.btn_process.pack(pady=10)

        self.panel_original = tk.Label(root)
        self.panel_original.pack(side="left", padx=20)

        self.panel_result = tk.Label(root)
        self.panel_result.pack(side="right", padx=20)

    def select_image(self):
        """
        打开文件对话框选择图像，并在左侧展示图像。

        :return: None
        """
        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.png;*.jpg;*.jpeg;*.bmp")])
        if file_path:
            self.image_path = file_path
            self.display_image(file_path, self.panel_original, size=(400, 400), title="原始图像")

    def process_image(self):
        """
        执行图像去噪处理，并在右侧展示处理后图像。

        :return: None
        """
        if not self.image_path:
            return

        denoised_img = self.run_denoise(self.image_path)
        if denoised_img is not None:
            temp_path = "./temp_result.png"
            cv2.imwrite(temp_path, denoised_img)
            self.display_image(temp_path, self.panel_result, size=(400, 400), title="处理后图像")

    def display_image(self, path, panel, size=(400, 400), title="图像"):
        """
        在指定标签控件中显示图像。

        :param path: 图像文件路径
        :param panel: 图像展示的 Label 控件
        :param size: 显示尺寸，默认 (400, 400)
        :param title: 显示标题文字
        :return: None
        """
        image = Image.open(path).convert("L")
        image = image.resize(size)
        image = ImageTk.PhotoImage(image)
        panel.configure(image=image, text=title, compound='top')
        panel.image = image

    def run_denoise(self, noisy_path):
        """
        执行图像恢复过程，包含非局部相似块搜索、SVD低秩逼近与锐化操作。

        :param noisy_path: 含噪图像路径（.png/.jpg 等）
        :return: 去噪并锐化后的图像（灰度图，np.uint8 类型）
        """
        img_true_path = noisy_path  # 默认真实图像路径与噪声图一致
        if not os.path.exists(img_true_path):
            return None

        _, _, img_denoised = denoise.main(
            block_size=8,
            step=2,
            n=32,
            k=4,
            method="manhattan",
            img_noisy=noisy_path,
            img_path=img_true_path
        )

        sharpen_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
        img_sharpened = cv2.filter2D(img_denoised, -1, sharpen_kernel)

        return img_sharpened.astype(np.uint8)


if __name__ == "__main__":
    root = tk.Tk()
    app = DenoiseApp(root)
    root.mainloop()
